import numpy as np
import pymc3 as pm
import matplotlib.pyplot as plt
import arviz as az
import pandas as pd
import DataPreparation
import MappingTabel
import pickle
import os
from sklearn.metrics import roc_curve, auc, precision_score, recall_score


def train():
    # 1. 数据预处理
    data = DataPreparation.preprocess()

    # 2. 特征选择
    DataPreparation.get_correlation_matrix(data)
    data.drop('Parch', axis=1, inplace=True)
    data.drop('SibSp', axis=1, inplace=True)
    data.drop('FamilyNum', axis=1, inplace=True)
    data.drop('Age', axis=1, inplace=True)

    print(data[:2][:])

    X, y = data.iloc[:, 2:], data.iloc[:, 1]
    (_, beta_len) = X.shape

    # 2. 定义贝叶斯模型
    with pm.Model() as model:
        # 先验分布
        alpha = pm.Normal('alpha', mu=0, sd=10)
        beta = pm.Normal('beta', mu=0, sd=2, shape=beta_len)

        # 似然函数
        theta = 1 / (1 + pm.math.exp(-(alpha + pm.math.dot(X, beta))))
        likelihood = pm.Bernoulli('survived', p=theta, observed=y)

        # MCMC采样
        trace = pm.sample(1000, return_inferencedata=True)
        with open(trace_path, 'wb') as f:
            pickle.dump(trace, f)
        # print(trace)
        az.plot_trace(trace)
        az.plot_posterior(trace)
        print(az.summary(trace))
        plt.show()
        return trace


# 预测某一乘客是否幸存
def predict_survival(sex, fare, pclass, name, cabin, embarked):
    pclass_map = {
        1: 0,
        2: 1,
        3: 2
    }

    cabin_map = {
        'A': 0,
        'B': 1,
        'C': 2,
        'D': 3,
        'E': 4,
        'F': 5,
        'G': 6,
        'T': 7,
        'U': 8
    }

    embarked_map = {
        'C': 0,
        'Q': 1,
        'S': 2,
        'U': 3
    }

    pclass_type = 3
    name_type = len(MappingTabel.TitleToNo)
    cabin_type = 9
    embarked_type = 4

    sex = DataPreparation.process_sex(sex)
    name = DataPreparation.process_name(name)
    cabin = DataPreparation.process_cabin(cabin)

    test_sample = [sex, fare]
    for i in range(pclass_type):
        if i == pclass_map[pclass]:
            test_sample.append(1)
        else:
            test_sample.append(0)

    for i in range(name_type):
        if i == name:
            test_sample.append(1)
        else:
            test_sample.append(0)

    for i in range(cabin_type):
        if i == cabin_map[cabin]:
            test_sample.append(1)
        else:
            test_sample.append(0)

    for i in range(embarked_type):
        if i == embarked_map[embarked]:
            test_sample.append(1)
        else:
            test_sample.append(0)

    # 计算生还概率
    p = 1 / (1 + pm.math.exp(-(alpha + np.dot(beta, test_sample))))
    return np.mean(p.eval())


def test_survival():
    datas = pd.read_csv('./titanic.csv')
    datas = DataPreparation.fill_nan(datas)

    pred = []
    ground_truce = []
    for _, data in datas.iterrows():
        # print(data)
        p = predict_survival(sex=data[4], fare=data[9], pclass=data[2], name=data[3], cabin=data[10], embarked=data[11])
        pred.append(1 if p > 0.5 else 0)
        ground_truce.append(data[1])

    # 计算 AUC-ROC
    fpr, tpr, thresholds = roc_curve(ground_truce, pred)
    roc_auc = auc(fpr, tpr)

    # 绘制 ROC 曲线
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC) Curve')
    plt.legend(loc="lower right")
    plt.show()

    # 计算精准率（Precision）和召回率（Recall）
    precision = precision_score(ground_truce, pred)
    recall = recall_score(ground_truce, pred)

    # 打印结果
    print(f"AUC-ROC: {roc_auc}")
    print(f"Precision: {precision}")
    print(f"Recall: {recall}")


if __name__ == "__main__":
    # 调用模型
    trace_path = './trace.pkl'
    if os.path.exists(trace_path):
        print('模型存在，调用模型')
        with open(trace_path, 'rb') as f:
            trace = pickle.load(f)
    else:
        print('模型不存在，开始训练模型')
        trace = train()

    # 测试模型
    alpha = trace.posterior['alpha']
    beta = trace.posterior['beta']
    test_survival()

    # survival_prob = predict_survival(sex="male", fare=71.2833, pclass=1, name='"Cumings, Mrs. John Bradley (Florence Briggs Thayer)"', cabin='C85', embarked="C")
    # print(f"该乘客的生还概率为: {survival_prob:.2f}")
